<template>
  <div id="app" @keydown.delete="handleDeleteKey" tabindex="0">
    <!-- Sidebar -->
    <div class="sidebar">
      <div class="controls">
        <div class="control-group">
          <h3>1. 上传图片</h3>
          <label for="file-upload" class="file-input-label">选择图片</label>
          <input id="file-upload" type="file" @change="handleImageUpload" multiple accept="image/*" />
        </div>

        <div class="control-group bordered">
          <h3>2. 定义类别</h3>
          <input type="text" v-model="newClassName" @keyup.enter="addClass" placeholder="输入类别名后按回车" class="class-input">
          <div class="class-tags">
            <span v-for="(name, index) in classNames" :key="index" class="tag" :style="{ backgroundColor: getColor(name) }">
              {{ name }}
              <button @click="removeClass(index)" class="remove-tag">&times;</button>
            </span>
          </div>
        </div>

        <div class="control-group">
            <h3>3. 开始标注</h3>
            <select v-model="selectedClass" class="class-select" :disabled="classNames.length === 0">
                <option v-if="classNames.length === 0" value="">请先定义一个类别</option>
                <option v-for="name in classNames" :key="name" :value="name">{{ name }}</option>
            </select>
            <div class="instructions">
                <p>拖拽鼠标以绘制标注框.</p>
                <p>单击标注框可以选中.</p>
                <p>拖动选中的框体或控制点进行调整.</p>
                <p>按 'Delete' 键删除所选框.</p>
            </div>
        </div>

        <div class="control-group">
            <h3>4. 导出数据</h3>
            <button @click="exportToYOLO" :disabled="images.length === 0">导出为 YOLOv11 ZIP</button>
        </div>
      </div>
    </div>

    <!-- Right Pane -->
    <div class="right-pane">
      <!-- Main Content (Canvas) -->
      <div class="main-content">
        <div class="canvas-container" ref="canvasContainer">
          <canvas ref="canvas" :style="{ cursor: canvasCursor }"></canvas>
        </div>
      </div>

      <!-- Thumbnail Bar -->
      <div class="thumbnail-bar" v-if="images.length > 0">
          <button @click="prevImage" class="thumb-nav-btn">&lt;</button>
          <div class="thumbnail-list-container" ref="thumbnailContainer">
              <div v-for="(image, index) in images"
                  :key="image.id"
                  :ref="el => { if (el) imageRefs[index] = el }"
                  class="thumbnail-item"
                  :class="{ active: activeImageIndex === index }"
                  @click="selectImage(index)">
                  <img :src="image.src" :alt="image.name" />
                  <span>{{ image.name }}</span>
              </div>
          </div>
          <button @click="nextImage" class="thumb-nav-btn">&gt;</button>
      </div>
    </div>
  </div>
</template>

<script setup>
import { ref, onMounted, computed, watch, nextTick } from 'vue';
import JSZip from 'jszip';

// --- STATE ---
const images = ref([]);
const activeImageIndex = ref(null);
const newClassName = ref('');
const classNames = ref(['人', '汽车', '卡车', '自行车']);
const selectedClass = ref('人');
const colorMap = ref({});

const canvas = ref(null);
let ctx = null;

// UI Refs
const thumbnailContainer = ref(null);
const imageRefs = ref([]);

// Interaction State
const isDrawing = ref(false);
const isDragging = ref(false);
const isResizing = ref(false);
const selectedAnnotationIndex = ref(null);
const dragHandle = ref(null);
const canvasCursor = ref('crosshair');
let startMousePos = { x: 0, y: 0 };

// --- COMPUTED ---
const activeImage = computed(() => images.value[activeImageIndex.value] || null);
const selectedAnnotation = computed(() => {
    if (activeImage.value && selectedAnnotationIndex.value !== null) {
        return activeImage.value.annotations[selectedAnnotationIndex.value];
    }
    return null;
});

// --- LIFECYCLE ---
onMounted(() => {
  ctx = canvas.value.getContext('2d');
  canvas.value.addEventListener('mousedown', handleMouseDown);
  canvas.value.addEventListener('mousemove', handleMouseMove);
  canvas.value.addEventListener('mouseup', handleMouseUp);
  canvas.value.addEventListener('mouseout', handleMouseOut);
  generateColorMap();
});

// --- IMAGE HANDLING ---
const handleImageUpload = (event) => {
  const files = Array.from(event.target.files);
  const newImages = files.map((file, index) => ({
    id: Date.now() + index,
    file,
    src: URL.createObjectURL(file),
    name: `img_${String(images.value.length + index + 1).padStart(4, '0')}.${file.name.split('.').pop()}`,
    originalWidth: 0, originalHeight: 0,
    annotations: [],
  }));
  images.value.push(...newImages);
  if (activeImageIndex.value === null) selectImage(0);
};

const selectImage = (index) => {
  if (index < 0 || index >= images.value.length) return;
  activeImageIndex.value = index;
  selectedAnnotationIndex.value = null;
};

const prevImage = () => {
    if (activeImageIndex.value > 0) {
        selectImage(activeImageIndex.value - 1);
    }
};

const nextImage = () => {
    if (activeImageIndex.value < images.value.length - 1) {
        selectImage(activeImageIndex.value + 1);
    }
};

const scrollToActiveThumbnail = () => {
    nextTick(() => {
        const activeThumb = imageRefs.value[activeImageIndex.value];
        if (activeThumb) {
            activeThumb.scrollIntoView({
                behavior: 'smooth',
                inline: 'center',
                block: 'nearest'
            });
        }
    });
};

// --- CANVAS & DRAWING ---
const redrawCanvas = () => {
    if (!activeImage.value || !ctx) {
        if(ctx && canvas.value) ctx.clearRect(0, 0, canvas.value.width, canvas.value.height);
        return;
    };
    const img = new Image();
    img.src = activeImage.value.src;
    img.onload = () => {
        if (canvas.value.width !== img.naturalWidth) canvas.value.width = img.naturalWidth;
        if (canvas.value.height !== img.naturalHeight) canvas.value.height = img.naturalHeight;
        
        ctx.clearRect(0, 0, canvas.value.width, canvas.value.height);
        ctx.drawImage(img, 0, 0);

        activeImage.value.annotations.forEach((ann, index) => {
            const isSelected = index === selectedAnnotationIndex.value;
            drawAnnotation(ann, isSelected);
        });
    };
};

const drawAnnotation = (ann, isSelected) => {
    const color = getColor(ann.className);
    ctx.strokeStyle = color;
    ctx.lineWidth = isSelected ? 3 : 2;
    ctx.strokeRect(ann.x, ann.y, ann.w, ann.h);

    ctx.fillStyle = color;
    ctx.font = 'bold 16px Arial';
    const label = `${ann.className}`;
    const textMetrics = ctx.measureText(label);
    const textWidth = textMetrics.width;
    const textHeight = 16; // Approximation
    
    const rectX = ann.x;
    const rectY = ann.y > textHeight + 5 ? ann.y - textHeight - 5 : ann.y + 5;

    ctx.fillStyle = 'rgba(0, 0, 0, 0.6)';
    ctx.fillRect(rectX, rectY, textWidth + 8, textHeight + 4);
    ctx.fillStyle = '#FFFFFF';
    ctx.fillText(label, rectX + 4, rectY + textHeight);

    if (isSelected) {
        drawHandles(ann);
    }
};

const drawHandles = (ann) => {
    const handles = getHandles(ann);
    ctx.fillStyle = '#FFFFFF';
    ctx.strokeStyle = '#000000';
    ctx.lineWidth = 1;
    for (const key in handles) {
        const handle = handles[key];
        ctx.fillRect(handle.x - 5, handle.y - 5, 10, 10);
        ctx.strokeRect(handle.x - 5, handle.y - 5, 10, 10);
    }
};

watch(activeImageIndex, () => {
    if (activeImage.value) {
        const img = new Image();
        img.src = activeImage.value.src;
        img.onload = () => {
            activeImage.value.originalWidth = img.naturalWidth;
            activeImage.value.originalHeight = img.naturalHeight;
            redrawCanvas();
            scrollToActiveThumbnail();
        };
    } else {
        redrawCanvas();
    }
}, { immediate: true });

// --- MOUSE EVENTS & INTERACTIONS ---
const getMousePos = (event) => {
    const rect = canvas.value.getBoundingClientRect();
    return {
        x: (event.clientX - rect.left) * (canvas.value.width / rect.width),
        y: (event.clientY - rect.top) * (canvas.value.height / rect.height)
    };
};

const getHandles = (ann) => ({
    topLeft: { x: ann.x, y: ann.y }, topRight: { x: ann.x + ann.w, y: ann.y },
    bottomLeft: { x: ann.x, y: ann.y + ann.h }, bottomRight: { x: ann.x + ann.w, y: ann.y + ann.h },
});

const getHandleAtPos = (pos, ann) => {
    const handles = getHandles(ann);
    for (const name in handles) {
        if (Math.abs(pos.x - handles[name].x) < 8 && Math.abs(pos.y - handles[name].y) < 8) return name;
    }
    return null;
};

const handleMouseDown = (event) => {
    if (!activeImage.value) return;
    startMousePos = getMousePos(event);

    if (selectedAnnotation.value) {
        const handle = getHandleAtPos(startMousePos, selectedAnnotation.value);
        if (handle) {
            isResizing.value = true;
            dragHandle.value = handle;
            return;
        }
    }

    for (let i = activeImage.value.annotations.length - 1; i >= 0; i--) {
        const ann = activeImage.value.annotations[i];
        if (startMousePos.x >= ann.x && startMousePos.x <= ann.x + ann.w &&
            startMousePos.y >= ann.y && startMousePos.y <= ann.y + ann.h) {
            isDragging.value = true;
            selectedAnnotationIndex.value = i;
            redrawCanvas();
            return;
        }
    }

    selectedAnnotationIndex.value = null;
    if (!selectedClass.value) {
        alert("请在开始标注前选择一个类别。");
        return;
    }
    isDrawing.value = true;
    const newAnn = { x: startMousePos.x, y: startMousePos.y, w: 0, h: 0, className: selectedClass.value };
    activeImage.value.annotations.push(newAnn);
    selectedAnnotationIndex.value = activeImage.value.annotations.length - 1;
};

const handleMouseMove = (event) => {
    if (!activeImage.value) return;
    const currentPos = getMousePos(event);
    const dx = currentPos.x - startMousePos.x;
    const dy = currentPos.y - startMousePos.y;

    if (isResizing.value && selectedAnnotation.value) {
        const ann = selectedAnnotation.value;
        switch (dragHandle.value) {
            case 'topLeft': ann.x += dx; ann.y += dy; ann.w -= dx; ann.h -= dy; break;
            case 'topRight': ann.w += dx; ann.y += dy; ann.h -= dy; break;
            case 'bottomLeft': ann.x += dx; ann.w -= dx; ann.h += dy; break;
            case 'bottomRight': ann.w += dx; ann.h += dy; break;
        }
        startMousePos = currentPos;
    } else if (isDragging.value && selectedAnnotation.value) {
        selectedAnnotation.value.x += dx;
        selectedAnnotation.value.y += dy;
        startMousePos = currentPos;
    } else if (isDrawing.value && selectedAnnotation.value) {
        selectedAnnotation.value.w = dx;
        selectedAnnotation.value.h = dy;
    }
    
    updateCursor(event);
    redrawCanvas();
};

const handleMouseUp = () => {
    if (isDrawing.value || isResizing.value) {
        const ann = selectedAnnotation.value;
        if (ann) {
            if (ann.w < 0) { ann.x += ann.w; ann.w = -ann.w; }
            if (ann.h < 0) { ann.y += ann.h; ann.h = -ann.h; }
            if (ann.w < 5 || ann.h < 5) {
                activeImage.value.annotations.splice(selectedAnnotationIndex.value, 1);
                selectedAnnotationIndex.value = null;
            }
        }
    }
    isDrawing.value = false; isDragging.value = false; isResizing.value = false;
    dragHandle.value = null;
    redrawCanvas();
};

const handleMouseOut = (event) => {
    if (isDrawing.value || isResizing.value || isDragging.value) handleMouseUp(event);
};

const handleDeleteKey = () => {
    if (selectedAnnotation.value) {
        activeImage.value.annotations.splice(selectedAnnotationIndex.value, 1);
        selectedAnnotationIndex.value = null;
        redrawCanvas();
    }
};

const updateCursor = (event) => {
    const pos = getMousePos(event);
    if (selectedAnnotation.value) {
        const handle = getHandleAtPos(pos, selectedAnnotation.value);
        if (handle) {
            canvasCursor.value = (handle === 'topLeft' || handle === 'bottomRight') ? 'nwse-resize' : 'nesw-resize';
            return;
        }
        if (pos.x >= selectedAnnotation.value.x && pos.x <= selectedAnnotation.value.x + selectedAnnotation.value.w &&
            pos.y >= selectedAnnotation.value.y && pos.y <= selectedAnnotation.value.y + selectedAnnotation.value.h) {
            canvasCursor.value = 'move';
            return;
        }
    }
    canvasCursor.value = 'crosshair';
};

// --- CLASS & COLOR MANAGEMENT ---
const generateColorMap = () => {
    classNames.value.forEach(name => {
        if (!colorMap.value[name]) {
            colorMap.value[name] = `hsl(${ (Object.keys(colorMap.value).length * 137.5) % 360 }, 70%, 50%)`;
        }
    });
};
const getColor = (className) => colorMap.value[className] || '#cccccc';

const addClass = () => {
  const name = newClassName.value.trim();
  if (name && !classNames.value.includes(name)) {
    classNames.value.push(name);
    newClassName.value = '';
    if (!selectedClass.value) selectedClass.value = name;
    generateColorMap();
  }
};

const removeClass = (index) => {
  const removedClass = classNames.value[index];
  classNames.value.splice(index, 1);
  if (selectedClass.value === removedClass) {
      selectedClass.value = classNames.value.length > 0 ? classNames.value[0] : '';
  }
  images.value.forEach(img => {
      img.annotations = img.annotations.filter(ann => ann.className !== removedClass);
  });
  redrawCanvas();
};

// --- EXPORT ---
const exportToYOLO = async () => {
    const annotatedImages = images.value.filter(img => img.annotations.length > 0);
    if (annotatedImages.length === 0) {
        return alert("没有标注可以导出。请至少标注一张图片。");
    }

    const zip = new JSZip();
    const shuffled = [...annotatedImages].sort(() => 0.5 - Math.random());
    const trainCount = Math.floor(shuffled.length * 0.7);
    const valCount = Math.floor(shuffled.length * 0.2);
    const sets = {
        train: shuffled.slice(0, trainCount),
        val: shuffled.slice(trainCount, trainCount + valCount),
        test: shuffled.slice(trainCount + valCount)
    };

    for (const setName in sets) {
        const setImages = sets[setName];
        if (setImages.length === 0) continue;
        const imgFolder = zip.folder(`dataset/images/${setName}`);
        const labelFolder = zip.folder(`dataset/labels/${setName}`);
        for (const image of setImages) {
            imgFolder.file(image.name, image.file);
            const yoloStrings = image.annotations.map(ann => {
                const classIndex = classNames.value.indexOf(ann.className);
                const { originalWidth: w, originalHeight: h } = image;
                const xCenter = (ann.x + ann.w / 2) / w;
                const yCenter = (ann.y + ann.h / 2) / h;
                const width = ann.w / w;
                const height = ann.h / h;
                return `${classIndex} ${xCenter.toFixed(6)} ${yCenter.toFixed(6)} ${width.toFixed(6)} ${height.toFixed(6)}`;
            });
            const labelFileName = image.name.substring(0, image.name.lastIndexOf('.')) + '.txt';
            labelFolder.file(labelFileName, yoloStrings.join('\n'));
        }
    }

    const yamlContent = `train: ../dataset/images/train/\nval: ../dataset/images/val/\ntest: ../dataset/images/test/\n\nnc: ${classNames.value.length}\nnames: [${classNames.value.map(name => `'${name}'`).join(', ')}]`;
    zip.file('dataset/data.yaml', yamlContent);

    const content = await zip.generateAsync({ type: 'blob' });
    const link = document.createElement('a');
    link.href = URL.createObjectURL(content);
    link.download = 'yolo11_dataset.zip';
    link.click();
    URL.revokeObjectURL(link.href);
};

</script>

<style>
:root {
  font-family: system-ui, Avenir, Helvetica, Arial, sans-serif;
  background-color: #242424;
  color: #e0e0e0;
}
body, html { margin: 0; padding: 0; height: 100%; overflow: hidden; }
#app { display: flex; height: 100vh; outline: none; }

.sidebar {
  width: 300px;
  flex-shrink: 0;
  background-color: #1e1e1e;
  padding: 1rem;
  display: flex;
  flex-direction: column;
  overflow-y: auto;
  border-right: 1px solid #333;
}

.right-pane {
  flex-grow: 1;
  display: flex;
  flex-direction: column;
  overflow: hidden;
}

.main-content {
  flex-grow: 1;
  display: flex;
  justify-content: center;
  align-items: center;
  background-color: #2d2d2d;
  overflow: auto;
  padding: 20px;
}
canvas {
  max-width: 100%;
  max-height: 100%;
  object-fit: contain;
  display: block;
  border-radius: 4px;
  box-shadow: 0 0 10px rgba(0,0,0,0.5);
}

.control-group { margin-bottom: 1.5rem; }
.control-group.bordered {
    border: 1px solid #444;
    padding: 1rem;
    border-radius: 8px;
    background-color: #2a2a2a;
}
h3 { margin-top: 0; margin-bottom: 0.75rem; color: #a0a0a0; font-size: 1rem; text-transform: uppercase; }
button, .file-input-label {
  width: 100%;
  padding: 0.8rem 1rem;
  font-size: 1rem;
  background-color: #4f46e5;
  color: white;
  border: none;
  border-radius: 5px;
  cursor: pointer;
  transition: background-color 0.2s;
  text-align: center;
}
button:hover, .file-input-label:hover { background-color: #4338ca; }
button:disabled { background-color: #555; cursor: not-allowed; }
input[type="file"] { display: none; }

.class-input, .class-select {
    width: 100%;
    box-sizing: border-box;
    padding: 0.6em;
    margin-bottom: 0.5rem;
    border-radius: 4px;
    border: 1px solid #555;
    background-color: #333;
    color: #fff;
}
.class-tags { display: flex; flex-wrap: wrap; gap: 0.5rem; margin-top: 0.5rem; }
.tag {
    color: white;
    padding: 0.4em 0.8em;
    border-radius: 12px;
    font-size: 0.9em;
    display: flex;
    align-items: center;
    font-weight: bold;
}
.remove-tag {
    background: none; border: none; color: white; cursor: pointer;
    margin-left: 0.5em; padding: 0; font-size: 1.3em; line-height: 1;
    opacity: 0.7;
}
.remove-tag:hover { opacity: 1; }

.instructions {
    font-size: 0.85rem;
    color: #888;
    margin-top: 1rem;
    background-color: #2a2a2a;
    padding: 0.5rem 0.8rem;
    border-radius: 4px;
}
.instructions p { margin: 0.3rem 0; }

.thumbnail-bar {
    flex-shrink: 0;
    background-color: #1a1a1a;
    border-top: 1px solid #333;
    display: flex;
    align-items: center;
    padding: 0.5rem;
    box-sizing: border-box;
    height: 120px; /* Fixed height for the bar */
}
.thumbnail-list-container {
    flex-grow: 1;
    display: flex;
    overflow-x: auto;
    gap: 0.5rem;
    scrollbar-width: none; /* Firefox */
}
.thumbnail-list-container::-webkit-scrollbar { display: none; /* Chrome, Safari, Opera */ }
.thumbnail-item {
    cursor: pointer;
    padding: 0.3rem;
    border: 2px solid transparent;
    border-radius: 6px;
    text-align: center;
    transition: all 0.2s;
    flex-shrink: 0; /* Prevent thumbnails from shrinking */
}
.thumbnail-item:hover { background-color: #444; }
.thumbnail-item.active { border-color: #4f46e5; background-color: #333; }
.thumbnail-item img {
    width: 80px;
    height: 80px;
    object-fit: cover;
    border-radius: 4px;
    display: block;
}
.thumbnail-item span {
    font-size: 0.75rem;
    color: #aaa;
    display: block;
    margin-top: 0.2rem;
    white-space: nowrap;
    overflow: hidden;
    text-overflow: ellipsis;
    width: 80px;
}
.thumb-nav-btn {
    width: 40px;
    height: 40px;
    margin: 0 0.5rem;
    flex-shrink: 0;
    font-size: 1.5rem;
    padding: 0;
    line-height: 40px;
}
</style>